module obs_profiler_mod

  use datetime
  use container
  use flogger
  use fiona
  use data_translators_utils_mod
  use latlon_process_mod
  use latlon_parallel_mod
  use latlon_process_mod
  use profiler_mod
  use profiler_zrada_mod
  use const_mod
  use namelist_mod
  use obs_operator_types_mod
  use obs_err_mod
  use efso_utils_mod

  implicit none

  private

  public obs_profiler_type
  public profiler_record_type

  type obs_profiler_type
    integer :: obs_size  = 0
    integer :: max_level = 0
    type(hash_table_type), allocatable :: platforms
    type(linked_list_type), allocatable :: records
    type(obs_grid_connector_type), allocatable :: connectors(:)
    type(profiler_record_type), allocatable :: simulated_mean_records(:)
    type(profiler_record_type), allocatable :: mean_simulated_records(:)
    type(profiler_record_type), allocatable :: simulated_records(:,:)
    type(profiler_record_type), allocatable :: innov_records(:)
    type(obs_err_type), allocatable :: obs_err(:)
    integer, allocatable :: mask(:)
  contains
    procedure :: init => obs_profiler_init
    procedure :: read => obs_profiler_read
    procedure :: clear => obs_profiler_clear
    procedure :: filter_by_latlon_and_time => obs_profiler_filter_by_latlon_and_time
    procedure :: calc_mean_simulated_records => obs_profiler_calc_mean_simulated_records
    procedure :: calc_perb_simulated_records => obs_profiler_calc_perb_simulated_records
    procedure :: calc_innov => obs_profiler_calc_innov
    final :: obs_profiler_final
  end type obs_profiler_type

contains

  subroutine obs_profiler_init(this)

    class(obs_profiler_type), intent(inout) :: this

    call this%clear()

    allocate(this%platforms)
    allocate(this%records  )

    this%platforms = hash_table(chunk_size=10000)

  end subroutine obs_profiler_init

  subroutine obs_profiler_read(this, file_path, mute)

    class(obs_profiler_type), intent(inout) :: this
    character(*), intent(in) :: file_path
    logical, intent(in) :: mute

    call profiler_zrada_read(file_path, this%platforms, this%records, mute)

  end subroutine obs_profiler_read

  subroutine obs_profiler_clear(this)

    class(obs_profiler_type), intent(inout) :: this

    if (allocated(this%platforms             )) deallocate(this%platforms             )
    if (allocated(this%records               )) deallocate(this%records               )
    if (allocated(this%connectors            )) deallocate(this%connectors            )
    if (allocated(this%simulated_mean_records)) deallocate(this%simulated_mean_records)
    if (allocated(this%mean_simulated_records)) deallocate(this%mean_simulated_records)
    if (allocated(this%simulated_records     )) deallocate(this%simulated_records     )
    if (allocated(this%innov_records         )) deallocate(this%innov_records         )
    if (allocated(this%obs_err               )) deallocate(this%obs_err               )
    if (allocated(this%mask                  )) deallocate(this%mask                  )

  end subroutine obs_profiler_clear

  subroutine obs_profiler_filter_by_latlon_and_time(this, start_lon, end_lon, start_lat, end_lat, start_time, end_time)

    class(obs_profiler_type), intent(inout) :: this
    real(r8), intent(in) :: start_lon
    real(r8), intent(in) :: end_lon
    real(r8), intent(in) :: start_lat
    real(r8), intent(in) :: end_lat
    type(datetime_type), intent(in) :: start_time
    type(datetime_type), intent(in) :: end_time

    type(linked_list_iterator_type) it
    integer n, i, j, k

    it = linked_list_iterator(this%records)
    do while (.not. it%ended())
      select type (record => it%value)
      type is (profiler_record_type)
        if (record%station%lon < start_lon * deg .or. record%station%lon > end_lon * deg .or. &
            record%station%lat < start_lat * deg .or. record%station%lat > end_lat * deg .or. &
            record%time < start_time .or. record%time > end_time) then
          call this%records%remove_item(it%item)
        end if
      end select
      call it%next()
    end do

    if (parallel_ensemble .and. proc%is_root_proc()) then
      call log_notice('After filter there is ' // to_str(this%records%size) // ' profiler records.')
    end if

    if (allocated(this%connectors)) deallocate(this%connectors)
    allocate(this%connectors(this%records%size))
    if (allocated(this%simulated_mean_records)) deallocate(this%simulated_mean_records)
    allocate(this%simulated_mean_records(this%records%size))
    if (allocated(this%mean_simulated_records)) deallocate(this%mean_simulated_records)
    allocate(this%mean_simulated_records(this%records%size))
    if (allocated(this%simulated_records)) deallocate(this%simulated_records)
    allocate(this%simulated_records(this%records%size,ens_ibeg:ens_iend))
    if (allocated(this%innov_records)) deallocate(this%innov_records)
    allocate(this%innov_records(this%records%size))

    n = 0
    it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => it%value)
      type is (profiler_record_type)
        n = n + record%pro%num_level; this%max_level  = max(this%max_level , record%pro%num_level)
      end select
      call it%next()
    end do
    this%obs_size = n

    ! Initialize observation error structure.
    allocate(this%obs_err(this%obs_size))
    n = 1
    it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => it%value)
      type is (profiler_record_type)
        do j = 1, record%pro%num_level
          this%obs_err(n)%u = 2 ! FIXME: Set obs_err for profiler
          this%obs_err(n)%v = 2
          n = n + 1
        end do
      end select
      call it%next()
    end do

    ! Set mask for missing values.
    allocate(this%mask(this%obs_size * 2))
    it = linked_list_iterator(this%records)
    i = 1
    do while (.not. it%ended())
      select type (record => it%value)
      type is (profiler_record_type)
        do k = 1, record%pro%num_level
          this%mask(i) = merge(0, 1, is_missing(record%pro%ua(k))); i = i + 1
          this%mask(i) = merge(0, 1, is_missing(record%pro%va(k))); i = i + 1
        end do
      end select
      call it%next()
    end do

  end subroutine obs_profiler_filter_by_latlon_and_time

  subroutine obs_profiler_calc_mean_simulated_records(this)

    class(obs_profiler_type), intent(inout) :: this

    real(r8), allocatable :: tmp(:,:)
    integer i, iens

    do i = 1, this%records%size
      this%mean_simulated_records(i)%pro%ua = 0
      this%mean_simulated_records(i)%pro%va = 0
    end do
    do iens = ens_ibeg, ens_iend
      do i = 1, this%records%size
        this%mean_simulated_records(i)%pro%ua = add(this%simulated_records(i,iens)%pro%ua, this%mean_simulated_records(i)%pro%ua)
        this%mean_simulated_records(i)%pro%va = add(this%simulated_records(i,iens)%pro%va, this%mean_simulated_records(i)%pro%va)
      end do
    end do

    ! Communicate with other processes.
    allocate(tmp(this%max_level,this%records%size))
    ! - u
    do i = 1, this%records%size
      tmp(:this%mean_simulated_records(i)%pro%num_level,i) = this%mean_simulated_records(i)%pro%ua
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%pro%ua = tmp(:this%mean_simulated_records(i)%pro%num_level,i) / ensemble_size
    end do
    ! - v
    do i = 1, this%records%size
      tmp(:this%mean_simulated_records(i)%pro%num_level,i) = this%mean_simulated_records(i)%pro%va
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%pro%va = tmp(:this%mean_simulated_records(i)%pro%num_level,i) / ensemble_size
    end do
    deallocate(tmp)

  end subroutine obs_profiler_calc_mean_simulated_records

  subroutine obs_profiler_calc_perb_simulated_records(this)

    class(obs_profiler_type), intent(inout) :: this

    type(linked_list_iterator_type) it
    integer i, j, k, n, iens

    ! Set observation error since pressure is interpolated when obs_operator_interp_run.
    if (proc%is_root_proc()) call log_notice('Set profiler observation error after pressure is interpolated.')
    n = 1
    it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => it%value)
      type is (profiler_record_type)
        do j = 1, record%pro%num_level
          if (is_missing(record%pro%p(j))) cycle
          do k = 1, size(profiler_err_plev)
            if (profiler_err_plev_bnds(1,k) <= record%pro%p(j) .and. record%pro%p(j) < profiler_err_plev_bnds(2,k)) exit
          end do
          if (k > size(profiler_err_plev)) call log_error('PROFILER does not found observation error for level ' // to_str(int(record%pro%p(j))) // 'hPa!')
          this%obs_err(n)%u = u_err_profiler(k)
          this%obs_err(n)%v = v_err_profiler(k)
          n = n + 1
        end do
      end select
      call it%next()
    end do

    ! Perturbation = Y - mean(Y)
    do iens = ens_ibeg, ens_iend
      do i = 1, this%records%size
        this%simulated_records(i,iens)%pro%ua = this%simulated_records(i,iens)%pro%ua - this%mean_simulated_records(i)%pro%ua
        this%simulated_records(i,iens)%pro%va = this%simulated_records(i,iens)%pro%va - this%mean_simulated_records(i)%pro%va
      end do
    end do

  end subroutine obs_profiler_calc_perb_simulated_records

  subroutine obs_profiler_calc_innov(this)

    class(obs_profiler_type), intent(inout) :: this

    type(linked_list_iterator_type) it
    integer i, j, k

    ! Innovation = Yo - H mean(X)
    it = linked_list_iterator(this%records)
    j = 0
    do i = 1, this%records%size
      select type (record => it%value)
      type is (profiler_record_type)
        call this%innov_records(i)%init_with_size(num_level =record%pro%num_level)
        this%innov_records(i)%pro%ua = record%pro%ua - this%simulated_mean_records(i)%pro%ua
        this%innov_records(i)%pro%va = record%pro%va - this%simulated_mean_records(i)%pro%va
        do k = 1, record%pro%num_level
          j = j + 1; call check_omb_ua(this%innov_records(i)%pro%ua(k), this%mask(j))
          j = j + 1; call check_omb_va(this%innov_records(i)%pro%va(k), this%mask(j))
        end do
      end select
      call it%next()
    end do

  end subroutine obs_profiler_calc_innov

  subroutine obs_profiler_final(this)

    type(obs_profiler_type), intent(inout) :: this

    call this%clear()

  end subroutine obs_profiler_final

end module obs_profiler_mod
